import torch.nn as nn
import torch

# vgg网络模型配置列表，数字表示卷积核个数，'M'表示最大池化层
cfgs = {
    'vgg11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],  # 模型A
    'vgg13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],  # 模型B
    'vgg16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],  # 模型D
    'vgg19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
    # 模型E
}


# 卷积层提取特征
def make_features(cfg: list):  # 传入的是具体某个模型的参数列表
    layers = []
    in_channels = 3  # 输入的原始图像(rgb三通道)
    for v in cfg:
        # 最大池化层
        if v == "M":
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        # 卷积层
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            layers += [conv2d, nn.ReLU(True)]
            in_channels = v
    # Sequential()方法是一个容器，描述了神经网络的网络结构，在Sequential()的输入参数中描述从输入层到输出层的网络结构
    return nn.Sequential(*layers)  # 单星号(*)将参数以元组(tuple)的形式导入


class VGG(nn.Module):
    # 初始化函数
    def __init__(self, model_type='vgg11', num_classes=1000, init_weights=False):
        super(VGG, self).__init__()
        self.features = make_features(cfgs[model_type])  # 卷积层提取特征
        # Sequential()方法是一个容器
        self.classifier = nn.Sequential(  # 全连接层进行分类
            # 每个训练批次中忽略一半的特征检测器，可以明显地减少过拟合现象
            nn.Dropout(p=0.5),
            # 全连接层
            nn.Linear(512 * 7 * 7, 2048),
            # 非线性激活函数 ReLU
            nn.ReLU(True),
            nn.Dropout(p=0.5),
            nn.Linear(2048, 2048),
            nn.ReLU(True),
            nn.Linear(2048, num_classes)
        )
        if init_weights:
            self._initialize_weights()

    # 前向传播
    def forward(self, x):
        # N x 3 x 224 x 224
        x = self.features(x)  # 运行特征提取模块（cfg数据字典里面的结构）
        # N x 512 x 7 x 7
        # flatten(dim)是对多维数据的降维函数，从第dim个维度开始展开，将后面的维度转化为一维.也就是说，只保留dim之前的维度，其他维度的数据全都挤在dim这一维。
        x = torch.flatten(x, start_dim=1)  # 展平降维
        # N x 512*7*7
        # 分类器 self.classifier
        x = self.classifier(x)
        x = nn.functional.log_softmax(x, dim=1)
        return x

    # 权重初始化
    def _initialize_weights(self):
        for m in self.modules():
            # isinstance 判断一个对象是否是一个已知的类型，类似 type()
            if isinstance(m, nn.Conv2d):
                # kaiming_normal_ 使用正态分布对输入张量进行赋值，深度学习还存在这样几种模型初始化策略：Constant 、Random 、Xavier
                # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                nn.init.xavier_uniform_(m.weight)  # Xavier均匀分布
                if m.bias is not None:
                    # constant(tensor, val)：用val的值填充输入的张量或变量
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                # normal(tensor, mean=0, std=1)：从给定均值mean和标准差std的正态分布N(mean, std)中生成值，填充输入的张量或变量
                # nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)
